{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "884f3125",
   "metadata": {},
   "source": [
    "# LoRA inference with NVIDIA NIM\n",
    "\n",
    "This is a demonstration of running inference against a LoRA adapter deployed with NVIDIA NIM. NIM supports LoRA adapters in .nemo (from NeMo Framework), and Hugging Face model formats. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b161f16b",
   "metadata": {
    "tags": []
   },
   "source": [
    "This notebook includes instructions to send an inference call to NVIDIA NIM using the Python `requests` library."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54759732",
   "metadata": {},
   "source": [
    "## Before you begin\n",
    "Ensure that you satisfy the pre-requisites, and have completed the setup instructions provided in the README associated with this tutorial to deploy the NIM container with LoRA."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e7917da",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa2477e9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import requests\n",
    "import json"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39d9918a",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Check available LoRA models\n",
    "\n",
    "Once the NIM server is up and running, check the available models as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2d71965",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "url = 'http://0.0.0.0:8000/v1/models'\n",
    "\n",
    "response = requests.get(url)\n",
    "data = response.json()\n",
    "\n",
    "print(json.dumps(data, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d5a93ac",
   "metadata": {},
   "source": [
    "This will return all the models available for inference by NIM. In this case, it will return the base model, as well as the LoRA adapters that were provided during NIM deployment - `llama3.1-8b-law-titlegen`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7c19acd",
   "metadata": {},
   "source": [
    "---\n",
    "## LoRA inference\n",
    "\n",
    "Inference can be performed by sending POST requests to the `/completions` endpoint.\n",
    "\n",
    "A few things to note:\n",
    "* The `model` parameter in the payload specifies the model that the request will be directed to. This can be the base model `meta/llama3.1-8b-instruct`, or any of the LoRA models, such as `llama3.1-8b-law-titlegen`.\n",
    "* `max_tokens` parameter specifies the maximum number of tokens to generate. At any point, the cumulative number of input prompt tokens and specified number of output tokens to generate should not exceed the model's maximum context limit. For llama3-8b-instruct, the context length supported is 8192 tokens.\n",
    "\n",
    "Following code snippets show how it's possible to send requests belonging to different LoRAs (or tasks). NIM dynamically loads the LoRA adapters and serves the requests. It also internally handles the batching of requests belonging to different LoRAs to allow better performance and more efficient of compute."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3edd2a9e",
   "metadata": {},
   "source": [
    "### Title Generation\n",
    "\n",
    "Try sending an example from the test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf6ea42a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "url = 'http://0.0.0.0:8000/v1/completions'\n",
    "\n",
    "headers = {\n",
    "    'accept': 'application/json',\n",
    "    'Content-Type': 'application/json'\n",
    "}\n",
    "\n",
    "# Example from the test set, following the template we trained the lora with\n",
    "prompt=\"Generate a concise, engaging title for the following legal question on an internet forum. The title should be legally relevant, capture key aspects of the issue, and entice readers to learn more. \\nQUESTION: In order to be sued in a particular jurisdiction, say New York, a company must have a minimal business presence in the jurisdiction. What constitutes such a presence? Suppose the company engaged a New York-based Plaintiff, and its representatives signed the contract with the Plaintiff in New York City. Does this satisfy the minimum presence rule? Suppose, instead, the plaintiff and contract signing were in New Jersey, but the company hired a law firm with offices in New York City. Does this qualify? \\nTITLE: \"\n",
    "data = {\n",
    "    \"model\": \"llama3.1-8b-law-titlegen\",\n",
    "    \"prompt\": prompt,\n",
    "    \"max_tokens\": 25,\n",
    "    \"temperature\":0\n",
    "}\n",
    "\n",
    "response = requests.post(url, headers=headers, json=data)\n",
    "response_data = response.json()\n",
    "\n",
    "print(json.dumps(response_data, indent=4))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
